'''
The files uploaded to telegram will have this naming convention:
[s_file]_[index]  example: 1_3128

The maximum filename length for a file is 64 ASCII chars (for telegram)
2 chars will be allocated for the session file part

This means that we have 10^62 filenames possible with only digits
(for each session file)

Don't upload files that are in the same directory as data_path
^ Don't remember why :0

Due to how files are downloaded, downloading 2 files with the same name
(not path) at the same time will cause problems.
^ Again i don't remember what i meant with this, is it a problem when you
download 2 files with different paths buth same last name or
is it a problem when the full path is the same?

Also when a file with same name and download path as one of previous files has
been downloaded then the original file will be replaced.
(If the original has not been moved)
'''

import logging
# Disable messages from pyrogram
logging.getLogger('pyrogram').setLevel(logging.CRITICAL)

from pyrogram import Client
from os import path, makedirs
from math import ceil
import asyncio

class TelegramHandler:
    def __init__(self,
                 config: dict,
                 s_file: int,
                 progress_fun: callable, # Pointer to progress function
                 data_fun: callable, # Called for multi chunk transfers
                 async_fun):

        try:
            self.telegram_channel_id = int(config['telegram']['channel_id'])
        except ValueError:
            self.telegram_channel_id = config['telegram']['channel_id']

        self.data_path = config['paths']['data_path']
        self.tmp_path = config['paths']['tmp_path']
        self.cipher_key = config['remote']['cipher_key']
        self.s_file = s_file # we need this for the naming when uploading
        self.progress_fun = progress_fun
        self.data_fun = data_fun
        self.async_files = async_fun
        self.download_full_path = config.getboolean('paths',
                                                    'download_full_path')
        self.should_stop = 0
        self.tot_chunks = 0

        self.mul_chunk_size = 1999
        self.chunk_size = self.mul_chunk_size << 20 # * (1024**2)

        self.telegram = Client(path.join(self.data_path, "a{}".format(s_file)),
                               config['telegram']['api_id'],
                               config['telegram']['api_hash'],
                               no_updates = True, sleep_threshold = -1)

        '''
        Connect to telegram servers when starting so that if we are missing any
        sessions it will prompt for login before starting the UI
        '''
        with self.telegram:
            pass


    async def uploadFiles(self, fileData: dict):
        self.tot_chunks = ceil(fileData['size'] / self.chunk_size)

        while True: # not end of file
            copied_file_path = path.join(self.tmp_path,
                                 "{}_{}".format(self.s_file, fileData['index']))

            fileData['chunkIndex'] = await self.async_files.splitFile(
                fileData['chunkIndex'], fileData['path'].encode(),
                copied_file_path.encode(),
                self.mul_chunk_size, 1024**2,
                self.cipher_key.encode()
            )

            async with self.telegram:
                msg_obj = await self.telegram.send_document(
                        self.telegram_channel_id, copied_file_path,
                        progress = self.progress_fun,
                        progress_args=(len(fileData['fileID']), self.tot_chunks,
                                       self.s_file))

            await self.async_files.remove(copied_file_path)
            # delete the chunk

            if self.should_stop == 2: # force stop
                if not self.isMultiChunk():
                    self.transferCleanup()
                    return
                break

            fileData['fileID'].append(msg_obj.id)
            fileData['index'] += 1

            if not fileData['chunkIndex']: # reached EOF
                break

            self.data_fun(fileData, self.s_file)

            if self.should_stop == 1:
                break

        self.transferCleanup()

        if not fileData['chunkIndex']: # Finished uploading
            return {'fileData' : {'rPath'    : fileData['rPath'],
                                  'fileID'   : fileData['fileID'],
                                  'size'     : fileData['size'],
                                  'checksum' : fileData['checksum']},
                    'index'    : fileData['index']}
            # Return file information


    async def downloadFiles(self, fileData: dict):
        finished_transfer = False
        self.tot_chunks = len(fileData['fileID'])

        if self.download_full_path:
            final_dir_path = (
                path.join(fileData['dPath'], *fileData['rPath'][:-1])
                    if fileData['dPath'] else
                path.join(self.data_path, "downloads", *fileData['rPath'][:-1]))

            if not path.isdir(final_dir_path):
                makedirs(final_dir_path)

            relative_file_path = path.join(*fileData['rPath'])
        else:
            relative_file_path = fileData['rPath'][-1]

        final_file_path = (
            path.join(fileData['dPath'], relative_file_path)
                if fileData['dPath'] else
            path.join(self.data_path, "downloads", relative_file_path))

        tmp_file_path = path.join(self.tmp_path,
                                  f"{fileData['rPath'][-1]}_chunk")

        # Begin downloading
        while True:
            async with self.telegram:
                message = await self.telegram.get_messages(
                                        self.telegram_channel_id,
                                        fileData['fileID'][fileData['IDindex']])

                await message.download(
                    file_name = tmp_file_path, progress = self.progress_fun,
                    progress_args = (fileData['IDindex'],
                                     self.tot_chunks, self.s_file)
                )

            # The download failed: [-503 Timeout]  upload.GetFile
            if not path.isfile(tmp_file_path):
                await asyncio.sleep(100)
                continue

            if self.should_stop == 2: # force stop
                break

            await self.async_files.concatFiles(
                tmp_file_path.encode(),
                final_file_path.encode(),
                1024 * 1024, self.cipher_key.encode()
            )
            await self.async_files.remove(tmp_file_path)

            fileData['IDindex'] += 1
            if fileData['IDindex'] == self.tot_chunks:
                # Finished or canceled with 1 but it was last chunk
                finished_transfer = True
                break

            self.data_fun(fileData, self.s_file)

            if self.should_stop == 1:
                # Issued normal cancel
                break

        self.transferCleanup()

        if finished_transfer:
            return final_file_path


    '''
    Mode is 1 for everything except IDList,
            2 for only IDList
    '''
    async def deleteUseless(self, IDList: list, mode: int = 1):
        deletedList = []

        async with self.telegram:
            if mode == 1:
                async for tFile in self.telegram.iter_history(
                                                      self.telegram_channel_id):
                    if (tFile.media) and (not tFile.id in IDList):
                        deletedList.append(tFile.id)

                if deletedList:
                    await self.telegram.delete_messages(
                        self.telegram_channel_id, deletedList
                    )

            elif mode == 2:
                await self.telegram.delete_messages(
                    self.telegram_channel_id, IDList
                )

        return deletedList


    '''
    Values of stop_type:
    1 - Wait until the current chunk transfer ended and appended
    2 - Cancel transfer, will still wait for appending to finish,
        not working currently
    '''
    async def stop(self, stop_type: int):
        if not stop_type in (1, 2):
            raise IndexError("stop_type should be 1 or 2.")
        if not self.isMultiChunk() and stop_type == 1:
            raise IndexError(
                "stop_type can't be 1 when transmitting single chunk files."
            )

        self.should_stop = stop_type
        if stop_type == 2: # force stop
            async with self.telegram:
                await self.telegram.stop_transmission()


    def isMultiChunk(self):
        if self.tot_chunks > 1:
            return True
        else:
            return False


    def transferCleanup(self):
        self.tot_chunks = 0
        self.should_stop = 0
